#include <cuda.h>

#if CUDA_VERSION >= 10010
#ifndef TRT_GELU_PLUGIN_H
#define TRT_GELU_PLUGIN_H

#include "NvInferPlugin.h"
#include "utils.h"
#include <string>
#include <vector>

extern "C"  int geluKernelLauncher(cudaStream_t stream,
                                const float* input, 
                                float* output, int len);

class GeLuPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
    GeLuPlugin() = delete;
    GeLuPlugin(const std::string name, const nvinfer1::DataType dataType);
    GeLuPlugin(const std::string name, const void* serialData, size_t serialLength);
public:
    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;

    // 算子输出张亮的维度
    nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex, 
                                        nvinfer1::DimsExprs const* inputs, 
                                        int32_t nbInputs, 
                                        nvinfer1::IExprBuilder& exprBuilder) noexcept override;

    bool supportsFormatCombination(int32_t pos, 
                                nvinfer1::PluginTensorDesc const* inOut, 
                                int32_t nbInputs, 
                                int32_t nbOutputs) noexcept override;
    
    void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, 
                        int32_t nbInputs,
                        nvinfer1::DynamicPluginTensorDesc const* out, 
                        int32_t nbOutputs) noexcept override;
    
    size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, 
                            int32_t nbInputs,
                            nvinfer1::PluginTensorDesc const* outputs, 
                            int32_t nbOutputs) const noexcept override;
    
    int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, 
                    nvinfer1::PluginTensorDesc const* outputDesc,
                    void const* const* inputs, 
                    void* const* outputs, 
                    void* workspace, 
                    cudaStream_t stream) noexcept override;

    nvinfer1::DataType getOutputDataType(int32_t index, 
                                        nvinfer1::DataType const* inputTypes, 
                                        int32_t nbInputs) const noexcept override;

    // IPluginV2 Methods
    char const* getPluginType()    const noexcept override;
    char const* getPluginVersion() const noexcept override;
    int32_t     getNbOutputs()     const noexcept override;
    
    int32_t initialize() noexcept override;
    void    terminate()  noexcept override;
    void    destroy()    noexcept override;
   
    size_t getSerializationSize()   const noexcept override;
    void serialize(void* buffer)    const noexcept override;
    
    void        setPluginNamespace(const char* pluginNamespace) noexcept override;
    char const* getPluginNamespace() const noexcept override;
private:
    const std::string  _pluginName;
    std::string        _namespace;
    nvinfer1::DataType _dataType;
    size_t             _ld;
};


class GeLuPluginCreater : public nvinfer1::IPluginCreator {
public:
    GeLuPluginCreater();

    const char* getPluginName()    const noexcept override;
    const char* getPluginVersion() const noexcept override;

    const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;

    nvinfer1::IPluginV2DynamicExt* createPlugin(const char* name, 
        const nvinfer1::PluginFieldCollection* fc) noexcept override;

    nvinfer1::IPluginV2DynamicExt* deserializePlugin(const char* name, 
        const void* serialData, size_t serialLength) noexcept override;

    void        setPluginNamespace(const char* pluginNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;
private:
    static nvinfer1::PluginFieldCollection _fc;          // 插件字段集合
    // 插件属性的集合，一个向量vector
    static std::vector<nvinfer1::PluginField> _pluginAttributes;   
    std::string _namespace;
};

#endif // TRT_GELU_PLUGIN_H
#endif // CUDA_VERSION >= 10010
